import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as signal
import scipy

def realignment_coeffs(data):
    """
    对系数进行隔列插零
    """
    z = np.zeros(data.shape[0])
    for i in range(data.shape[1]-1,0,-1):
        data = np.insert(data,i,z,axis=1)
    return data

def realignment_data(data,channel_num):
    disp_len = int(np.ceil(data.size / channel_num))
    patch_size = int(disp_len * channel_num - data.size)
    patch_data = np.concatenate((data, np.zeros(patch_size)))
    polyphase_data = np.zeros(patch_data.size*2) 
    half = (channel_num // 2)
    for i in range(patch_data.size // half):
        if(i == (patch_data.size // half - 1)):
            polyphase_data[i*channel_num+half:(i+1)*channel_num] = patch_data[i*half:(i+1)*half]
        else:
            polyphase_data[i*channel_num+half:(i+1)*channel_num+half] = list(patch_data[i*half:(i+1)*half])*2
    polyphase_data = polyphase_data.reshape( (channel_num, -1), order='F')
    polyphase_data = np.flip(polyphase_data,0)
    return polyphase_data

def gen_filter_coeffs(numtaps, M):
    coeffs = scipy.signal.firwin(numtaps*M, cutoff=2.0/M, window="hamming")
    coeffs = np.reshape(coeffs, (M, -1), order='F')
    coeffs = realignment_coeffs(coeffs)
    return coeffs

def polyphase_filter(data,filter_coeffs,channel_num):   
     
    polyphase_data = realignment_data(data, channel_num)
    polyphase_data = polyphase_data.reshape( (channel_num, -1), order='F')
    
    # print(polyphase_data)
    
    filt_data = np.zeros(polyphase_data.shape)
    for k in range(channel_num):
        filt_data[k] = scipy.signal.lfilter(filter_coeffs[k], 1, polyphase_data[k])
        # filt_data[k] = scipy.signal.rfilter(filter_coeffs[k], 1, filt_data[k])

    dispatch_data = scipy.fft.ifft(filt_data, axis=0)
    return dispatch_data

def apply_chirp(freq,size,bw=400.0,fc=1382.0):
    dm = 2.64476
    dm_dispersion = 2.41e-4
    dispersion_per_MHz = 1e6 * dm / dm_dispersion
    phasors = np.zeros((size),dtype=complex)
    binwidth = -bw / size
    coeff = 2 * np.pi * dispersion_per_MHz / (fc * fc);
    for i in range (size):
        f = i * binwidth + 0.5 * bw
        phasors[i] = np.exp(1j * coeff * f * f / (fc+f));
        if i == 0:
            phasors[i] = 0; 
        freq[i] = phasors[i] * freq[i];

# def delay_dispersion(freq1,freq2):
    
      

def get_period_size(bw):
    period = 0.00575730363767324
    return int(period * bw * 1e6)

def fold_data(data,block_size,psize):
    cur = 0
    pdata = np.zeros((psize))
    pnum = np.zeros((psize))
    for i in range(block_size):
        if(cur >= psize):
            cur = 0
        pnum[cur] = pnum[cur] + 1
        pdata[cur] = (pnum[cur] - 1) * pdata[cur] / pnum[cur] + data[i] / pnum[cur]
        cur = cur + 1
    return pdata
            
def integral_data(data,size,n=1024):
    data = np.abs(data)
    t = size // n
    d = np.zeros((n),dtype=np.float64)
    sum = 0
    j = 0
    z = 0
    for i in range(size):
        if j == t:
            d[z] = (sum / t)
            sum = 0
            j = 0
            z = z + 1
        j = j+1
        sum += data[i]
    return d

def load_time(filename):
    data = np.loadtxt(filename,dtype=np.float64)
    n = data.shape[0]
    return data,n

def load_freq(filename):
    data = np.loadtxt(filename)
    n = data.shape[0] // 2
    freq = np.zeros((n),dtype=complex)
    for i in range (n):
        freq[i] = data[2*i] + data[2*i+1]*1j
    return freq,n

def plot(data):
    plt.figure(figsize=(20,5),dpi=100)
    plt.ylabel("Magnitude(dB)")
    plt.xlabel("Phase")
    plt.plot(np.abs(data))
    plt.show()

def time_data(filename):
    data,n = load_time(filename)
    freq = scipy.fft.fft(data)
    apply_chirp(freq,n//2)
    idata = np.abs(scipy.fft.ifft(freq[:n//2]))
    psize = get_period_size(400)
    pdata = fold_data(idata,n//2,psize)
    idata = integral_data(pdata,psize)
    # idata = integral_data(idata,n//2)
    plot(idata[:1000])

def freq_data(filename):
    freq,n = load_freq(filename)
    apply_chirp(freq,n)
    idata = np.abs(scipy.fft.ifft(freq))
    psize = get_period_size(400)
    pdata = fold_data(idata,n,psize)
    idata = integral_data(pdata,psize)
    plot(idata)
    
def opfb_data(filename,channel_num,M):
    data,n = load_time(filename)
    coeffs = gen_filter_coeffs(M,channel_num)    
    freq = polyphase_filter(data,coeffs,channel_num)
    freq_size = freq.shape[1]
    N = int(freq_size*4)
    myfreq = np.zeros((N), dtype=complex)
    start = 0
    end = 0
    mystart = 0
    myend = 0
    bw = 0
    for i in range(9):
        mystart += bw
        if i == 0:
            start = freq_size // 2
            end = freq_size // 2 + freq_size // 4
        elif i == 8:
            start = freq_size // 2 - freq_size // 4
            end = freq_size // 2 
        else:
            start = freq_size // 2 - freq_size // 4
            end = freq_size // 2 + freq_size // 4
        bw = end - start
        myend = mystart + bw
        if i % 2 == 0:
            myfreq[mystart:myend] = scipy.fft.fftshift(scipy.fft.fft(freq[i]))[start:end]
        else:
            myfreq[mystart:myend] = scipy.fft.fft(freq[i])[start:end]
    
    apply_chirp(myfreq,N)
    df = np.abs(scipy.fft.ifft(myfreq))
    psize = get_period_size(400)
    pdata = fold_data(df,N,psize)
    idata = integral_data(pdata,psize)
    # idata = integral_data(df,N)
    np.savetxt("myopfb.txt",idata)
    plot(idata[:1000])
    # return freq,myfreq,N,freq_size

time_data("time_26.txt") # 直接使用时域数据，消色散画图

# freq_data("freq_22.txt") # 使用频域数据，消色散画图

# opfb_data("time_26.txt",16,64)

# tt,myfreq,N,n = opfb_data("time_26.txt",16,64)

# plot(myfreq)

# s = n // 4
# e = 3 * n // 4
# n = (n // 2)

# freq = scipy.fft.fft(tt[2])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1482)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test2.txt",idata);
# # # # idata = integral_data(fdata,n)
# # plot(idata)

# freq = scipy.fft.fft(tt[3])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1432)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test3.txt",idata);

# freq = scipy.fft.fft(tt[4])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1382)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test4.txt",idata);

# freq = scipy.fft.fft(tt[5])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1332)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test5.txt",idata);

# freq = scipy.fft.fft(tt[6])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1282)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test6.txt",idata);

# freq = scipy.fft.fft(tt[7])
# freq = freq[s:e]
# apply_chirp(freq,n,50,1232)
# fdata = np.abs(scipy.fft.ifft(freq))
# psize = get_period_size(50)
# pdata = fold_data(fdata,n,psize)
# idata = integral_data(pdata,psize)
# np.savetxt("test6.txt",idata);
    
# 1582 1569.5 1557 
# 1557 1532.0 1507
# 1507 1482.0 1457
# 1457 1432.0 1407
# 1407 1382.0 1357
# 1357 1332.0 1307
# 1307 1282.0 1257
# 1257 1232.0 1207
# 1207 1194.5 1182
# start = 1582
# end = 0
# fc = 1582.0
# bw = 0
# for i in range(9):
#     fc = fc - bw / 2
#     if i == 0 or i == 8:
#         bw = 25
#         end = start - bw
#     else:
#         bw = 50
#         end = start - bw
#     fc = fc - bw / 2
#     print(start,fc,end)
#     start  = start - bw
